# 引入依赖库
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn import preprocessing
from models import *
from datetime import datetime, timedelta
from joblib import load

# 参数定义
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
#input_length = 28  # 输入时间窗口长度（过去4周）
# output_length = 7  # 输出预测长度（未来1周）
features_num = 11  # 特征数量（flow + 3个时间特征）
# dim = 256  # 模型隐藏层维度
# num_blocks = 2  # LSTM层数
scalar = True  # 必须与训练时相同
scalar_contain_labels = True


# 加载预训练模型
def load_pretrained_model(model_path,params):
    # 定义模型结构（必须与训练时一致）
    model = LSTMMainWithPadding(
        input_size=features_num,
        output_len=params["output_length"],
        lstm_hidden=params["dim"],
        lstm_layers=params["num_blocks"],
        batch_size=1,  # 预测时batch_size=1
        device=device
    )
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.to(device)
    model.eval()  # 切换到评估模式
    return model


# 数据预处理函数（必须与训练时逻辑一致）
def preprocess_data(df,holidays,stationInfo,scalerPath):
    # 添加时间特征
    df['day_of_week'] = df['date'].dt.dayofweek
    df['day_of_month'] = df['date'].dt.day
    df['month'] = df['date'].dt.month

    holidays['date'] = pd.to_datetime(holidays['date'])
    holiday_cols = ['date', 'is_holiday', 'holiday_type', 'holiday_num', 'serial_num', 'is_free']
    df = df.merge(holidays[holiday_cols], on='date', how='left')
    stationInfo['date'] = pd.to_datetime(stationInfo['date'])
    stationInfo_cols = ['date', 'station_status', 'station_hour']
    df = df.merge(stationInfo[stationInfo_cols], on='date', how='left')
    # 填充空值（非节假日）
    df['is_holiday'] = df['is_holiday'].fillna(0)
    df['holiday_type'] = df['holiday_type'].fillna(0)
    df['holiday_num'] = df['holiday_num'].fillna(0)
    df['serial_num'] = df['serial_num'].fillna(0)
    df['is_free'] = df['is_free'].fillna(0)
    df['station_status'] = df['station_status'].fillna(0)
    df['station_hour'] = df['station_hour'].fillna(0)


    # 提取特征并归一化
    features_ = df[
        ['flow', 'day_of_week', 'day_of_month', 'month', 'is_holiday', 'holiday_type', 'holiday_num', "serial_num",
         "is_free","station_status","station_hour"]].values
    #features_ = df[['flow', 'day_of_week', 'day_of_month', 'month']].values

    # 注意：此处应加载训练时的scaler，假设已保存为scaler.pkl
    # 此处简化处理，实际应保存scaler避免重新拟合
    scaler = load(scalerPath)

    features_norm = scaler.transform(features_)
    return features_norm, scaler
    # scaler = preprocessing.MinMaxScaler()

    #
    # if scalar:
    #     features_norm = scaler.transform(features_)
    #     if scalar_contain_labels:
    #         labels_norm = features_norm[:, 0]
    #     else:
    #         labels_norm = df['flow'].values
    # else:
    #     features_norm = features_
    #     labels_norm = df['flow'].values




# 预测函数
def predict_future(model, scaler, current_date, df,param,holidays,scalerPath):
    # 1. 读取数据并筛选最近input_length天的数据
    df['date'] = pd.to_datetime(df['date'])
    df = df.sort_values('date').reset_index(drop=True)

    # 获取当前日期对应的最近input_length天数据
    end_date = current_date
    start_date = end_date - timedelta(days=param["input_length"])
    mask = (df['date'] >= start_date) & (df['date'] < end_date)
    recent_data = df.loc[mask].copy()

    print(recent_data)
    # 检查数据长度是否足够
    # if len(recent_data) < input_length:
    #     raise ValueError(f"数据不足{input_length}天！当前仅{len(recent_data)}天")

    # 2. 预处理数据
    features_norm, _ = preprocess_data(recent_data,holidays,scalerPath)

    # 3. 构建输入序列（直接取最后input_length天）
    input_sequence = features_norm[-param["input_length"]:]

    # 4. 预测
    with torch.no_grad():
        input_tensor = torch.FloatTensor(input_sequence).unsqueeze(0).to(device)  # 添加batch维度
        prediction_norm = model(input_tensor).cpu().numpy()[0]

    # 5. 逆归一化
    if scalar_contain_labels and scalar:
        temp_pred = np.zeros((1, param["output_length"], features_num))
        temp_pred[0, :, 0] = prediction_norm
        prediction = scaler.inverse_transform(temp_pred[0])[:, 0]
    else:
        prediction = prediction_norm

    # 6. 生成预测日期
    pred_dates = [end_date + timedelta(days=i-1) for i in range(1, param["output_length"] + 1)]
    # 改为这样可能更准确
    #pred_dates = [end_date + timedelta(days=i) for i in range(param["output_length"])]
    return prediction, pred_dates

